# -*- coding: utf-8 -*-
# -*- coding: utf-8 -*-

from pymongo import UpdateOne, DESCENDING
from factor.base_factor import BaseFactor
from data.finance_report_crawler import FinanceReportCrawler
from data.data_module import DataModule
from util.stock_util import get_all_codes, get_trading_dates
from util.database import DB_CONN
import pandas as pd
import numpy as np
from util.erio_util import dualdate, dualdateb, dual_ammount,dualbasicdate,dualcode

"""
实现市盈率因子的计算和保存
"""


class SP_TTM_Factor(BaseFactor):
    def __init__(self):
        BaseFactor.__init__(self, name='pe')

    def compute(self, begin_date, end_date):
        """
        计算指定时间段内所有股票的该因子的值，并保存到数据库中
        :param begin_date:  开始时间
        :param end_date: 结束时间
        """
        dm = DataModule()
        dates = get_trading_dates(begin_date, end_date)
        codes = ['601808']
        update_requests = []
        for code in codes:
            data = dm.get_stock_zcfzb(code)
            basicdata = dm.get_stock_basic(dualcode(code))
            E2Pall = dm.get_stock_zcfzb(code)
            for date in dates:
                for asa in data['report_date'].apply(dualdate):
                    if date < asa:
                        continue
                    else:
                        reportdat = asa
                        break
                E2P = E2Pall[E2Pall['report_date'] == dualdateb(reportdat)]['SUMLASSET'].astype('float').tolist()[0]

                TTM = basicdata[basicdata['trade_date'] == dualbasicdate(date)]['total_mv'].tolist()[0]*10000
                E2P_TTM = (E2P / TTM)

                update_requests.append(
                    UpdateOne(
                        {'code': code, 'date': date},
                        {'$set': {'code': code, 'date': date, 'E2P_TTM': E2P_TTM}}, upsert=True))
                print(update_requests)
                break
            break


SP_TTM_Factor().compute('2018-01-02', '2018-09-05')
